#include "session/session.h"

#include "smart_serve_log.h"
#include "service/service.h"

namespace OHOS {
namespace SmartServe {

void MergeIDs(const unsigned* ids1, unsigned n1, const unsigned* ids2, unsigned n2, unsigned* dst, unsigned& ndst)
{
    std::vector<int> all;
    all.reserve(n1 + n2);

    all.insert(all.end(), ids1, ids1 + n1);
    all.insert(all.end(), ids2, ids2 + n2);

    std::sort(all.begin(), all.end());
    auto end = std::unique(all.begin(), all.end());

    std::copy(all.begin(), end, dst);
    ndst = static_cast<unsigned>(end - all.begin());
}

SessionResConfig Session::conservativeResConfig() const
{
    if (!pending_) {
        return currentResConfig_;
    }

    auto ret = SessionResConfig{
        .memSize_ = std::max(pendingResConfig_.memSize_, currentResConfig_.memSize_),
    };

#define LMS_MERGE_IDS(ids, n)                                                                                          \
    MergeIDs(currentResConfig_.ids, currentResConfig_.n, pendingResConfig_.ids, pendingResConfig_.n, ret.ids, ret.n);

    LMS_MERGE_IDS(cpuIDs, nCPUs);
    LMS_MERGE_IDS(gpuIDs, nGPUs);
    LMS_MERGE_IDS(npuIDs, nNPUs);

#undef LMS_MERGE_IDS

    return ret;
}

bool Session::isRunnable() const
{
    auto& rc = (pending_) ? pendingResConfig_ : currentResConfig_;

    for (auto& cfg : resConfigs_) {
        if (cfg.isSatisfiedBy(rc)) {
            return true;
        }
    }

    return false;
}

bool Session::isCurrentlyRunnable() const
{
    auto& rc = currentResConfig_;

    for (auto& cfg : resConfigs_) {
        if (cfg.isSatisfiedBy(rc)) {
            return true;
        }
    }

    return false;
}

void Session::ThreadMain()
{
    while (true) {
        {
            std::unique_lock<std::mutex> lock(mutex_);
            cv_.wait(lock, [this]() { return pending_ || terminating_; });
        }

        if (pending_) {
            HandlePending();
        } // slots: (t) + (p)

        if (terminating_) { // slots: t + (p)
            if (pending_) { // slots: t + p
                HandlePending();
            } // slots: t
            break;
        } // slots: (t) + (p)

        // 工作循环
        HandleTask();

    } // slots: t

    HandleFinalize();
}

void Session::TriggerPending(const SessionResConfig& newResConfig)
{
    pendingResConfig_ = newResConfig;
    pending_ = true;

    cv_.notify_all(); // TODO：目前持有锁的情况下被调用
}

void Session::HandlePending()
{
    SMART_SERVE_LOGD("[Session] HandlePending");

    ASSERT(pending_);

    SessionApplyResult result = {ApplyResult::UNCHANGED, {}};
    if (!terminating_) {
        result = adapter_->ApplyResConfig(pendingResConfig_);
    }
    FinishPendingUse(result);
}

void Session::FinishPendingUse(const SessionApplyResult& applyResult)
{
    auto result = applyResult.result_;
    bool successful = (result == ApplyResult::SUCCESSFUL);

    // 生成回复 Event
    auto event = std::make_unique<SessionAppliedEvent>(currentResConfig_, pendingResConfig_, successful);

    // 处理 Fatal
    if (result == ApplyResult::FATAL) {
        std::unique_lock<std::mutex> lock(Service::Instance()->msMgmtMutex_);
        TriggerTerminating();
    }

    {
        std::unique_lock<std::mutex> lock(mutex_);

        // 修改 resConfig
        if (successful) {
            currentResConfig_ = pendingResConfig_;
        }

        if (!applyResult.newResConfigs_.empty()) {
            resConfigs_ = std::move(applyResult.newResConfigs_);
        }

        // 清理 pending_
        pending_ = false;
    }

    // 通知 Scheduler
    Service::Instance()->scheduler_->Schedule(std::move(event));
}

void Session::HandleTask()
{
    bool runnable = isCurrentlyRunnable();
    while (runnable) {
        SMART_SERVE_LOGD("[Session] HandleTask");

        // 获取任务
        std::shared_ptr<EvalTask> task;
        {
            std::unique_lock<std::mutex> lock(mutex_);
            if (tasks_.empty()) {
                break;
            }
            task = tasks_.front();
        }

        if (task->state_ == EvalTask::TaskState::CANCELLED) {
            std::unique_lock<std::mutex> lock(mutex_);
            tasks_.pop_front();
            break;
        }
        // 做任务
        auto evalResult = adapter_->Eval(task.get());

        // 处理任务结果
        switch (evalResult.result_) {
            case EvalResult::FINISHED: { // 任务结束
                std::unique_lock<std::mutex> lock(mutex_);
                tasks_.pop_front();
                break;
            }
            case EvalResult::CONTINUE: { // 任务单步完成
                // 无事可做
                break;
            }
            case EvalResult::FATAL:
            default: { // 不可恢复错误
                std::unique_lock<std::mutex> lock(Service::Instance()->msMgmtMutex_);
                TriggerTerminating();
                break;
            }
        }

        // 更新 resConfigs_
        if (!evalResult.newResConfigs_.empty()) {
            {
                std::unique_lock<std::mutex> lock(mutex_);
                resConfigs_ = std::move(evalResult.newResConfigs_);
            }

            runnable = isCurrentlyRunnable();
            if (!runnable) {
                Service::Instance()->scheduler_->Schedule(std::make_unique<ScheduleEvent>());
            }
        }

        // 检查异步事件
        if (pending_ || terminating_) {
            break;
        }
    }
}

void Session::HandleFinalize()
{
    SMART_SERVE_LOGD("[Session] HandleFinalize");

    // Abort 所有排队中任务
    AbortAllTasks();

    // 调用引擎回收资源
    adapter_->Finalize();
}

void Session::AbortAllTasks()
{
    std::unique_lock<std::mutex> lock(mutex_);
    for (auto& task : tasks_) {
        task->Done(EvalTask::TaskState::ABORTED);
    }
}

void Session::TriggerTerminating()
{
    SMART_SERVE_LOGD("[Session] TriggerTerminating");

    auto service = Service::Instance();

    // 移除管理结构
    {
        // 移除 Session
        auto self = service->sessionManager_.RemoveSession(this).release();

        // 检查是否已被他人终止
        if (self == nullptr) { // 因为维护管理结构由终止发起者完成，为空说明已经被终止过
            ASSERT(terminating_);
            return;
        }

        // 设置 Terminating
        AsyncTerminate();
        ASSERT(self == this);

        for (auto model : models_) {
            model->RemoveSession(this);
        }
    }

    // 通知 Maintainer 回收
    service->maintainer_.AddTask([this]() {
        WaitForTerminated();
        Service::Instance()->scheduler_->Schedule(std::make_unique<SessionReleseEvent>(this));
    });
}

} // namespace SmartServe
} // namespace OHOS